"""
ParentDocumentRetriever 是 LangChain 提供的一种文档检索策略，
主要用于解决 分块检索 时因文本分割导致的 上下文丢失问题。
它通过将文档拆分为小块存储，但在检索时返回原始的大文档（父文档），从而保留完整的上下文信息
"""

from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import InMemoryStore
from langchain_community.vectorstores import Chroma

from models import get_ollama_embeddings_client

embeddings = get_ollama_embeddings_client()

# 加载文档
loader = TextLoader("../data/document/state_of_the_union.txt", encoding="utf-8")
documents = loader.load()

# 定义父文档分割器（将文档分割成较大的块）
parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
# 定义子文档分割器（将父文档分割成较小的块用于存储）
child_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=50)

# 初始化向量存储
vectorstore = Chroma(
    collection_name="split_parents",
    embedding_function=embeddings,
)

# 初始化文档存储
docstore = InMemoryStore()

# 创建ParentDocumentRetriever实例
retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=docstore,
    child_splitter=child_splitter,
    parent_splitter=parent_splitter,
)

# 添加文档到检索器
retriever.add_documents(documents)

# 检索示例
query = "What did the president say about Ketanji Brown Jackson?"
retrieved_docs = retriever.get_relevant_documents(query)

# 打印结果
print(f"检索到 {len(retrieved_docs)} 个文档")
for i, doc in enumerate(retrieved_docs):
    print(f"\n文档 {i+1}:")
    print(f"内容: {doc.page_content[:500]}...")
    print(f"元数据: {doc.metadata}")
